#!/usr/bin/env python
# Licensed under the MIT license:
# http://www.opensource.org/licenses/mit-license.php 

"""
Class to use the preprocessor to get tweet words by user (as a dict, user=key).
"""

import csv
import json
import sys
import preprocessor

INPUT_FILENAME = 't_s_1009'  # default combined output from TweetObjectProcessor
OUTPUT_FILENAME = 'wordlist.csv'


class Wordlist:
  """Generates wordlist based on input, returns or writes to disk."""
  def __init__(self, input=INPUT_FILENAME, raw=None, skipread=False):
    if not skipread:
      file = open(input)
      self.text = file.readlines()
      file.close()
    elif raw:
      self.text = raw  # intended for CGI use, when we're not opening a file
    self.worddict = {}
    self.tweets = []
    self.pre = preprocessor.Preprocessor()

  def Process(self):
    """Process the input file to dict - don't call if skipread set to True."""
    for line in self.text:  # scraped tweets are single-line but that's okay
      try:
        objs = json.loads(line)
      except:
        continue
      for tweet in objs:  # may only be one but that's okay too
        if tweet.has_key('user') and tweet.has_key('text'):
          self.AddToDict(tweet['user']['screen_name'], tweet['text'])

  def AddToDict(self, user, text):
    """Add a given tweet to the dict - generally not for direct calling."""
    text = self.pre.Preprocess(text)
    self.tweets.append(' '.join(text))  # store the cleaned tweet
    if self.worddict.has_key(user):
      text.extend(self.worddict[user])
    self.worddict[user] = text

  def Save(self):
    """Save the dict to disk in csv format: user,words that they tweet."""
    outputfile = open(OUTPUT_FILENAME, 'w')
    csvwriter = csv.writer(outputfile)
    for row in self.worddict.items():
      row = list(row)
      row[1] = ' '.join(row[1])
      csvwriter.writerow(row)
    outputfile.close()
  
  def Load(self):
    """Load the dict from disk in csv format: user,words that they tweet."""
    inputfile = open(OUTPUT_FILENAME)
    csvreader = csv.reader(inputfile)
    for row in csvreader:
      self.worddict[row[0]] = row[1].split(' ')
